{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5bf4d107-2356-46ad-b668-25e4b6bf0c2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.autograd import Variable"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "712771a1",
   "metadata": {},
   "source": [
    "ConvGRU 是 GRU 的空间变体，专为处理视频帧或动态图像数据而设计。\n",
    "ConvGRU 是 GRU 的一种变体，但它的应用目标是**空间-时间序列数据**（例如视频帧序列）。 GRU 是一种通用的循环神经网络，用于处理时间序列数据。\n",
    "\n",
    "ConvGRU 输入的是二维或三维张量（例如图像帧或视频序列）  $x_t$ 。\n",
    "\n",
    "ConvGRU 中的线性变换使用的是**卷积操作**，即用卷积核（而不是矩阵乘法）来计算。\n",
    "\n",
    "ConvGRU：处理视频帧、动态场景分割等涉及空间和时间维度的数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f5bf341c-fa3c-47b4-957b-194feec88438",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvGRUCell(nn.Module):\n",
    "    def __init__(self, input_size, input_dim, hidden_dim, kernel_size, bias, dtype):\n",
    "        \"\"\"\n",
    "        Initialize the ConvLSTM cell\n",
    "        :param input_size: (int, int)\n",
    "            Height and width of input tensor as (height, width).\n",
    "        :param input_dim: int\n",
    "            Number of channels of input tensor.\n",
    "        :param hidden_dim: int\n",
    "            Number of channels of hidden state.\n",
    "        :param kernel_size: (int, int)\n",
    "            Size of the convolutional kernel.\n",
    "        :param bias: bool\n",
    "            Whether or not to add the bias.\n",
    "        :param dtype: torch.cuda.FloatTensor or torch.FloatTensor\n",
    "            Whether or not to use cuda.\n",
    "        \"\"\"\n",
    "        super(ConvGRUCell, self).__init__()\n",
    "        self.height, self.width = input_size\n",
    "        self.padding = kernel_size[0] // 2, kernel_size[1] // 2\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.bias = bias\n",
    "        self.dtype = dtype\n",
    "\n",
    "        self.conv_gates = nn.Conv2d(\n",
    "            in_channels=input_dim + hidden_dim,\n",
    "            out_channels=2 * self.hidden_dim,  # for update_gate,reset_gate respectively\n",
    "            kernel_size=kernel_size,\n",
    "            padding=self.padding,\n",
    "            bias=self.bias,\n",
    "        )\n",
    "\n",
    "        self.conv_can = nn.Conv2d(\n",
    "            in_channels=input_dim + hidden_dim,\n",
    "            out_channels=self.hidden_dim,  # for candidate neural memory\n",
    "            kernel_size=kernel_size,\n",
    "            padding=self.padding,\n",
    "            bias=self.bias,\n",
    "        )\n",
    "\n",
    "    def init_hidden(self, batch_size):\n",
    "        return Variable(\n",
    "            torch.zeros(batch_size, self.hidden_dim, self.height, self.width)\n",
    "        ).type(self.dtype)\n",
    "\n",
    "    def forward(self, input_tensor, h_cur):\n",
    "        \"\"\"\n",
    "        :param self:\n",
    "        :param input_tensor: (b, c, h, w)\n",
    "            input is actually the target_model\n",
    "        :param h_cur: (b, c_hidden, h, w)\n",
    "            current hidden and cell states respectively\n",
    "        :return: h_next,\n",
    "            next hidden state\n",
    "        \"\"\"\n",
    "        combined = torch.cat([input_tensor, h_cur], dim=1)\n",
    "        combined_conv = self.conv_gates(combined)\n",
    "\n",
    "        gamma, beta = torch.split(combined_conv, self.hidden_dim, dim=1)\n",
    "        reset_gate = torch.sigmoid(gamma)\n",
    "        update_gate = torch.sigmoid(beta)\n",
    "\n",
    "        combined = torch.cat([input_tensor, reset_gate * h_cur], dim=1)\n",
    "        cc_cnm = self.conv_can(combined)\n",
    "        cnm = torch.tanh(cc_cnm)\n",
    "\n",
    "        h_next = (1 - update_gate) * h_cur + update_gate * cnm\n",
    "        return h_next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d118e19c-57d7-4c8d-baf7-58de4061b4ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvGRU(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        input_dim,\n",
    "        hidden_dim,\n",
    "        kernel_size,\n",
    "        num_layers,\n",
    "        dtype,\n",
    "        batch_first=True,\n",
    "        bias=True,\n",
    "        return_all_layers=False,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        :param input_size: (int, int)\n",
    "            Height and width of input tensor as (height, width).\n",
    "        :param input_dim: int e.g. 256\n",
    "            Number of channels of input tensor.\n",
    "        :param hidden_dim: int e.g. 1024\n",
    "            Number of channels of hidden state.\n",
    "        :param kernel_size: (int, int)\n",
    "            Size of the convolutional kernel.\n",
    "        :param num_layers: int\n",
    "            Number of ConvLSTM layers\n",
    "        :param dtype: torch.cuda.FloatTensor or torch.FloatTensor\n",
    "            Whether or not to use cuda.\n",
    "        :param alexnet_path: str\n",
    "            pretrained alexnet parameters\n",
    "        :param batch_first: bool\n",
    "            if the first position of array is batch or not\n",
    "        :param bias: bool\n",
    "            Whether or not to add the bias.\n",
    "        :param return_all_layers: bool\n",
    "            if return hidden and cell states for all layers\n",
    "        \"\"\"\n",
    "        super(ConvGRU, self).__init__()\n",
    "\n",
    "        # Make sure that both `kernel_size` and `hidden_dim` are lists having len == num_layers\n",
    "        kernel_size = self._extend_for_multilayer(kernel_size, num_layers)\n",
    "        hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers)\n",
    "        if not len(kernel_size) == len(hidden_dim) == num_layers:\n",
    "            raise ValueError(\"Inconsistent list length.\")\n",
    "\n",
    "        self.height, self.width = input_size\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.kernel_size = kernel_size\n",
    "        self.dtype = dtype\n",
    "        self.num_layers = num_layers\n",
    "        self.batch_first = batch_first\n",
    "        self.bias = bias\n",
    "        self.return_all_layers = return_all_layers\n",
    "\n",
    "        cell_list = []\n",
    "        for i in range(0, self.num_layers):\n",
    "            cur_input_dim = input_dim if i == 0 else hidden_dim[i - 1]\n",
    "            cell_list.append(\n",
    "                ConvGRUCell(\n",
    "                    input_size=(self.height, self.width),\n",
    "                    input_dim=cur_input_dim,\n",
    "                    hidden_dim=self.hidden_dim[i],\n",
    "                    kernel_size=self.kernel_size[i],\n",
    "                    bias=self.bias,\n",
    "                    dtype=self.dtype,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        # convert python list to pytorch module\n",
    "        self.cell_list = nn.ModuleList(cell_list)\n",
    "\n",
    "    def forward(self, input_tensor, hidden_state=None):\n",
    "        \"\"\"\n",
    "        :param input_tensor: (b, t, c, h, w) or (t,b,c,h,w) depends on if batch first or not\n",
    "            extracted features from alexnet\n",
    "        :param hidden_state:\n",
    "        :return: layer_output_list, last_state_list\n",
    "        \"\"\"\n",
    "        if not self.batch_first:\n",
    "            # (t, b, c, h, w) -> (b, t, c, h, w)\n",
    "            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)\n",
    "\n",
    "        # Implement stateful ConvLSTM\n",
    "        if hidden_state is not None:\n",
    "            raise NotImplementedError()\n",
    "        else:\n",
    "            hidden_state = self._init_hidden(batch_size=input_tensor.size(0))\n",
    "\n",
    "        layer_output_list = []\n",
    "        last_state_list = []\n",
    "\n",
    "        seq_len = input_tensor.size(1)\n",
    "        cur_layer_input = input_tensor\n",
    "\n",
    "        for layer_idx in range(self.num_layers):\n",
    "            h = hidden_state[layer_idx]\n",
    "            output_inner = []\n",
    "            for t in range(seq_len):\n",
    "                # input current hidden and cell state then compute the next hidden and cell state through ConvLSTMCell forward function\n",
    "                h = self.cell_list[layer_idx](\n",
    "                    input_tensor=cur_layer_input[:, t, :, :, :], h_cur=h  # (b,t,c,h,w)\n",
    "                )\n",
    "                output_inner.append(h)\n",
    "\n",
    "            layer_output = torch.stack(output_inner, dim=1)\n",
    "            cur_layer_input = layer_output\n",
    "\n",
    "            layer_output_list.append(layer_output)\n",
    "            last_state_list.append([h])\n",
    "\n",
    "        if not self.return_all_layers:\n",
    "            layer_output_list = layer_output_list[-1:]\n",
    "            last_state_list = last_state_list[-1:]\n",
    "\n",
    "        return layer_output_list, last_state_list\n",
    "\n",
    "    def _init_hidden(self, batch_size):\n",
    "        init_states = []\n",
    "        for i in range(self.num_layers):\n",
    "            init_states.append(self.cell_list[i].init_hidden(batch_size))\n",
    "        return init_states\n",
    "\n",
    "    @staticmethod\n",
    "    def _check_kernel_size_consistency(kernel_size):\n",
    "        if not (\n",
    "            isinstance(kernel_size, tuple)\n",
    "            or (\n",
    "                isinstance(kernel_size, list)\n",
    "                and all([isinstance(elem, tuple) for elem in kernel_size])\n",
    "            )\n",
    "        ):\n",
    "            raise ValueError(\"`kernel_size` must be tuple or list of tuples\")\n",
    "\n",
    "    @staticmethod\n",
    "    def _extend_for_multilayer(param, num_layers):\n",
    "        if not isinstance(param, list):\n",
    "            param = [param] * num_layers\n",
    "        return param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a7aeb8fa-b41e-4af2-8ca4-879b8fc29ddb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1, 64, 6, 6])\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    # set CUDA device\n",
    "    os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "    # detect if CUDA is available or not\n",
    "    use_gpu = torch.cuda.is_available()\n",
    "    if use_gpu:\n",
    "        dtype = torch.cuda.FloatTensor # computation in GPU\n",
    "    else:\n",
    "        dtype = torch.FloatTensor\n",
    "\n",
    "    height = width = 6\n",
    "    channels = 256\n",
    "    hidden_dim = [32, 64]\n",
    "    kernel_size = (3,3) # kernel size for two stacked hidden layer\n",
    "    num_layers = 2 # number of stacked hidden layer\n",
    "    model = ConvGRU(input_size=(height, width),\n",
    "                    input_dim=channels,\n",
    "                    hidden_dim=hidden_dim,\n",
    "                    kernel_size=kernel_size,\n",
    "                    num_layers=num_layers,\n",
    "                    dtype=dtype,\n",
    "                    batch_first=True,\n",
    "                    bias = True,\n",
    "                    return_all_layers = False).cuda()\n",
    "\n",
    "    batch_size = 1\n",
    "    time_steps = 1\n",
    "    input_tensor = torch.rand(batch_size, time_steps, channels, height, width).cuda()  # (b,t,c,h,w)\n",
    "    layer_output_list, last_state_list = model(input_tensor)\n",
    "    print(layer_output_list[0].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "588b888b-8ac1-4b93-8fea-0b4f6513cf75",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
